import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
import csv
from DataSetLoader import *
from model import *

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

rnn_model = LSTMClassifier(5120, 360, 4, 4)
rnn_model.load_state_dict(torch.load('bestNetv2.pth'))
rnn_model.to(device)

rnn_model.eval()

LableEncode = {"N":0, "A":1, "O":2, "~":3 }
LableDecode = ["N", "A", "O", "~"]

dataSetDir = "./DataSetProcess/MyFFTValidation2017/"

with torch.no_grad():
    with open(dataSetDir + "REFERENCE.csv") as csvfile:
        with open('result.csv', 'w', newline='\n') as savefile:
            csv_reader = csv.reader(csvfile)  # 使用csv.reader读取csvfile中的文件
            csv_writer = csv.writer(savefile)
            pbar = tqdm(list(csv_reader))
            for i in pbar:
                filename = i[0]+".npy"
                data = np.load(dataSetDir + filename).astype(np.float32)
                data = torch.tensor(data).unsqueeze(0).to(device)
                pbar.set_description('Processing ' + filename)
                outputs = rnn_model(data)
                _, predicted = torch.max(outputs.data, 1)
                csv_writer.writerow([i[0], LableEncode[i[1]], int(predicted)])